#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/3/21
# @Author  : geekhch
# @Email   : geekhch@qq.com
# @Desc    : 项目运行入口

"""
Usage:
    run.py ner [options]
    run.py nre [options]
    run.py predict [options]

Options:
    -h --help                               show this screen.
    --cuda                                  use GPU
    --debug                                 debug mode
    --vocab=<file>                          vocab file
    --seed=<int>                            seed [default: 0]
    --epoch=<int>                           epoch size [default: 3]
    --batch-size=<int>                      batch size [default: 64]
    --embed-size=<int>                      embedding size [default: 256]
    --hidden-size=<int>                     hidden size [default: 1000]
    --clip-grad=<float>                     gradient clipping [default: 5.0]
    --lr-decay=<float>                      learning rate decay [default: 0.8]
    --lr-patience=<int>                     learning rate decay per iters [default: 300]
    --lr=<float>                            learning rate [default: 0.001]
    --log=<str>                             log description [default: normal]
"""

from docopt import docopt
import torch
from torch.optim.lr_scheduler import StepLR


from torch.utils.data import DataLoader
from utils import logger, device
from myPath import *
import os
from myNER.myDataset import CLUE

args = docopt(__doc__)
None if args['--debug'] else logger.level('INFO')

log_dir = DULE_NER_LOG

from datetime import datetime
if args['ner']:
    logger.add(f'{log_dir}/{datetime.now().strftime("%b%d_%H-%M-%S")}_{args["--log"]}.log')
if args['nre']:
    if not path.exists(DULE_LOG):
        os.makedirs(DULE_LOG)
    logger.add(f'{DULE_LOG}/{datetime.now().strftime("%b%d_%H-%M-%S")}_{args["--log"]}.log')

logger.info(str(args))

def train_ner():
    from myNER.net import BiLSTM_CRF
    from myNER.myDataset import CLUE, DUIE
    # train_loader = DataLoader(CLUE(path.join(CLUENER_DIR, 'train.json')),
    #                         collate_fn=CLUE.collate_fn, batch_size=int(args['--batch-size']),
    #                         num_workers=3, shuffle=True)
    # dev_loader = DataLoader(CLUE(path.join(CLUENER_DIR, 'dev.json')),
    #                         collate_fn=CLUE.collate_fn, batch_size=30,
    #                         num_workers=1, shuffle=True)
    train_loader = DUIE.get_loader('train', int(args['--batch-size']))
    dev_loader = DUIE.get_loader('test')

    # train_ner(model, train_loader, dev_loader)
    # model = BiLSTM_CRF(int(args['--hidden-size']), dataClass=DUIE)
    model = torch.load(f'{DULE_NER_MODEL}/{max(os.listdir(DULE_NER_MODEL))}', map_location='cpu')
    model = model.to(device)
    BiLSTM_CRF.evaluate(model, dev_loader, None, DUIE)
    # BiLSTM_CRF.train_lstm_ccrf(model, train_loader, dev_loader, args,
    #                           visual_log_dir=DULE_NER_VISUAL,
    #                           model_save_dir=DULE_NER_MODEL,
    #                           dataClass=DUIE)

def train_nre():
    from myRE.baidu_cnn_softmax import get_framework
    framework = get_framework()
    framework.train_model()

def predict(sentence):
    from myNER.net import BiLSTM_CRF
    best_ner = '04161335-1600-0.31_best(0.750).pt'
    best_nre = 'Apr16_00-10-06_cnn_best_89.9.pth.tar'

    ner_model = torch.load(f'{CLUENER_MODEL}/{best_ner}', map_location='cpu')
    from myRE.baidu_cnn_softmax import model as nre_model
    nre_model.load_state_dict(torch.load(f'{DULE_MODEL}/{best_nre}', map_location='cpu')['state_dict'])
    # print(ner_model)
    # print(nre_model)
    print('text:', sentence)
    print('***实体列表***')
    entities = BiLSTM_CRF.predict(ner_model, sentence, True, dataClass=CLUE)
    for k,v in entities.items():
        print(f'{k:<10s} {v[0]:<10s} {v[1]}')


    print('***关系列表***')
    token = list(sentence)
    for h in entities.keys():
        for t in entities.keys():
            if h == t:
                continue
                # "h": {"name": "顾漫", "pos": [53, 55], "type": "人物"},
            head = {'name':h, 'pos':entities[h][1], 'type':entities[h][0]}
            tail = {'name':t, 'pos':entities[t][1], 'type':entities[t][0]}
            sample = {'token':token, 'h':head, 't':tail}
            result = nre_model.infer(sample)
            if result[0] != '其他' and result[1] > 0.8:
                print(f'{h:<10s} {t:<10s} {result}')



if __name__ == '__main__':
    if args['ner']:
        train_ner()
    elif args['nre']:
        train_nre()
    elif args['predict']:
        sentence1 = '据新闻联播,根据世卫组织最新统计数据,13日全球新冠肺炎新增81577例,全球累计确诊新冠肺炎超过417万例,死亡28.7万例。'
        sentence2 = '网络大电影《大鱼》由北京新片场传媒股份有限公司独家宣发'
        sentence3 = '网络大电影《大鱼》由殷悦执导、杨春毅担任制片人，梅洋、付梦妮、狄龙、常戎新老演员联合出演，今日在爱奇艺强势上线。'
        sentence4 = '新浪娱乐讯 5月14日，应采儿Jasper一起拍摄的封面大片释出。大片中，应采儿挺着二胎孕肚出镜母爱满满，Jasper抱着妈妈，母子两人对着镜头甜笑画面超温馨。对于二胎家庭不可避免的“争宠”问题，应采儿说：“没人会计较爱是多或少，因为一家人是要互相关爱的，不分年龄大小。性格不同，表达爱的方式也不一样。'
        sentence5 = '新浪娱乐讯 近日，阿娇与赖弘国闪电离婚，男方微博取消关注阿娇，而阿Sa也被网友发现取消了对赖弘国的微博关注，以行动力挺姐妹。蔡卓妍（阿Sa），华语女歌手、演员，中国香港女子演唱组合Twins成员之一。'
        sentence6 = '当恽代英在泸州被军阀赖心辉扣押后，吴玉章立即去电保释，并聘请他到校担任教育学教员，还让他在礼堂给学生讲《阶级斗争》。'
        predict(sentence6)
        